"""
@Time: 2021-7-23
加强堆
"""


# 自定义加强堆
class heap_plus():
    def __init__(self, cmp):
        self.heap_list = []
        self.heap_index_map = {}  # 反向索引字典
        self.size = 0  # 堆长度
        self.cmp = cmp  # 自定义比较器函数

    def push(self, item):
        self.heap_list.append(item)
        # 新元素先添加到最后
        self.heap_index_map[self.size] = item
        self.heap_insert(self.size)
        self.size += 1

    def heap_insert(self, i):
        """
        从给定的索引处向上调整heap
        i为索引值, i的左孩子是2*i+1, 右孩子是2*i+2, 父节点是(i-1)/2向下取整
        :param i:
        :return:
        """
        index_dad = (i - 1) // 2
        while cmp(self.heap_list[i], self.heap_list[index_dad]) < 0:  # 小根堆
            self.swap(i, index_dad)
            i = index_dad
            index_dad = (i - 1) // 2

    def print(self):
        for i, item in enumerate(self.heap_list):
            print(i, ": ", item)

    # 列表和反向索引同时交换两个元素
    def swap(self, i, j):
        # if max(i, j) > self.size - 1:
        #     return None
        # if self.size < 2:
        #     return None
        # 交换列表
        self.heap_list[i], self.heap_list[j] = self.heap_list[j], self.heap_list[i]
        # 交换反向索引
        i = self.heap_index_map[i]
        j = self.heap_index_map[j]
        self.heap_index_map[j] = i
        self.heap_index_map[i] = j


if __name__ == '__main__':
    def cmp(o1, o2):
        return o1 - o2


    heap1 = heap_plus(cmp)
    heap1.push(5)
    heap1.push(1)
    heap1.push(9)
    heap1.print()
